from flask import Flask, request, jsonify
import logging, json
import requests
from db import db
import re
client = db()
import time

app = Flask(__name__)

 # 设置日志级别和格式
app.logger.setLevel(logging.DEBUG)
handler = logging.StreamHandler()
handler.setLevel(logging.DEBUG)
formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
handler.setFormatter(formatter)
app.logger.addHandler(handler)

file_handler = logging.handlers.TimedRotatingFileHandler(
	filename='logs/server.log',
	when='midnight',
	interval=1,
	backupCount=7,
)
app.logger.addHandler(file_handler)

# 假设这是你的安全令牌，实际应用中应该从配置文件或环境变量中获取
SECURITY_TOKEN = "sk-32hrhfyg3b8rghjefj09ekf0ki0fj9ujrg95jsdfa435evfhjuyu"
 
@app.before_request
def check_sk():
   # 获取请求头中的 SK 令牌
	sk = request.headers.get('SK')

# 检查 SK 是否存在且与预期值匹配
	if not sk or sk != SECURITY_TOKEN:
	# 如果验证失败，返回 401 Unauthorized 错误
		return jsonify({"error": "Unauthorized"}), 401
 
def get_params():
   # 检查是否有 POST 请求，并获取表单数据
	if request.method == 'POST':
		# 尝试获取 JSON 数据
		post_json = request.get_json(silent=True)
		if post_json:
			post_params = post_json
		else:
			# 如果不是 JSON，尝试获取表单数据
		   post_params = request.form.to_dict()
	elif request.method == 'GET':
		post_params = request.args.to_dict()

	return post_params

keys = {
	"task_id":         0,
	"prompt":          1,
	"negative_prompt": 2,
	"steps":           3,
	"sampler":         4,
	"face_fix":        5, 
	"cfg_scale":       6,
	"seed":            7,
	"height":          8,
	"width":           9,
	"hd_fix":          10,
	"hd_redraw_rate":  11, 
	"hd_scale":        12, 
	"hd_scale_alg":    13, 
	"hd_sample_num":   14, 
	"model":   15,
	"image":   16
}

@app.route('/run/predict', methods=['POST'])
def predict():
	params = get_params()
	app.logger.info(f'predict create :{params}')
	inputs = {key:value for key, value in zip(keys, params.get('data'))}
	if not inputs['image']:
		inputs.pop('image') ## 空值pop

	body = dict()
	body['version'] = client.get_version(inputs['model'])
	body['input'] = inputs


	app.logger.info(f'模型实际入参: {body}')
	if not body['version']:
		return '模型异常'

	## 发出请求
	url = 'https://api.replicate.com/v1/predictions'
	headers = {"Authorization":"Token r8_ZOYCNsQbdNu8Kox9uKSx9Rh3C6BMoVD0joeXf", "Content-Type":"application/json"}

	res = requests.post(url, headers=headers, data=json.dumps(body))
	resdata = json.loads(res.content)
	app.logger.info(f' 返回 : {resdata}')



	## 更新数据库，关联taskid和progressId
	task_id = params.get('data')[0]
	progress_id = resdata.get("id")
	app.logger.info(f'connect task_id: {task_id} ==> progress_id {progress_id}')
	client.update(task_id, progress_id)

	c = 0
	while(True):
		## 获取进度
		progress = get_progress(task_id)
		if progress == -1.0:
			return "failed"
		app.logger.info(f'progress : {progress}')
		time.sleep(1)
		c += 1
		
		if isinstance(progress, dict):
			return_format = {"Data":[[{'Name':progress.get('output')[0].split('/')[-1], 'IsFile':True, 'Data':""}], {"url":progress.get('output')[0], "seed":progress.get('seed')}], "IsGenerating":True, "Duration":0.0, "AverageDuration":0.0}

			return jsonify(return_format)
		
		
		if c > 28:
			return "failed"
 
def get_progress(task_id):
	def get_res(progress_id):
		url = f'https://api.replicate.com/v1/predictions/{progress_id}'
		headers = {"Authorization":"Token r8_ZOYCNsQbdNu8Kox9uKSx9Rh3C6BMoVD0joeXf", "Content-Type":"application/json"}

		## 获取task对应的progressid
		res = requests.get(url, headers=headers)
		res = json.loads(res.content)
		app.logger.info(f'进程查看:{res}')

		return res


	progress_id = client.get_progress(task_id)
	if not progress_id:
		return -1.0
		app.logger.error('获取progress_id失败')
	app.logger.info(f'get progress_id {progress_id} from {task_id}')

	res = get_res(progress_id)
	## 开始处理后播报
	nt = 1
	while(res.get('status') == 'starting'):
		time.sleep(nt)
		nt += 1
		res = get_res(progress_id)

	progress = re.findall('\d+%', res.get('logs', ''))
	current_progress = [0.01*float(i.replace('%',  '')) for i in progress] + [0.0]
	fseed = re.findall('seed:\s*(-?\d+)', res.get('logs', ''))
	seed = fseed[0] if fseed else -1

	if res.get('output'):
		## 生成文件后，获取并上传oss
		tmp = list()
		for item in res.get('output'):
			tmp.append(client.uploadfile(item))
		res['output'] = tmp

	## 队列中，未开始
	if res.get('status') == 'starting' and not res.get('logs'):
		time.sleep(1)

	## 失败直接返回
	if res.get('status') == 'failed':
		return -1
	
	if max(current_progress) == 1.0:
		client.update_image(task_id, res.get('output')[0])
		return {'output':res.get('output'), 'seed':seed}
	else:
		return max(current_progress)



@app.route('/internal/progress', methods=['POST'])
def progress():
	params = get_params()
	app.logger.info(f'check progress :{params}')
	## 根据获取到的task查找处理进程
	task_id = params.get('id_task')

	## 匹配格式
	v = get_progress(task_id)
	return_format = {
		"active": True,
		"queued": False,
		"completed": True,
		"progress": v if isinstance(v, float) else 1.0,
		"eta": 8.03231692314148,
		"live_preview": None,
		"id_live_preview": 1,
		"ImageName": "" if isinstance(v, float) else v.get('output')[0]  ,
		"textinfo": ""
	}

	return jsonify(return_format)
 
 
if __name__ == '__main__':
   # 运行 Flask 应用
	app.run('127.0.0.1', port=7890, debug=True)

	
